//! \file linux_arm_exidx_elf.c

#define _GNU_SOURCE

#include <link.h>

#include "linux_elf_utils.h"
#include "memory_utils.h"
#include "linux_arm_exidx_elf.h"

//------------------------------------------------------------------------------

#define EXIDX_CANTUNWIND 0x1
#define BIT_31_MASK (1 << 31)
#define EXIDX_COMPACT(word) ((word) & BIT_31_MASK)

//------------------------------------------------------------------------------

TBGError AddrFromPrel31(addr_t prel31, addr_t *ret_addr) {
    TBGError ret = 0;
    CHECK_CALL(ReadData(prel31, ret_addr, 1), exit_point, ret);
    CHECK_RET(ret_addr, NULL_PREL31, exit_point, ret);

    (*ret_addr) &= ~BIT_31_MASK; // remove the first bit
    (*ret_addr) += prel31; 
    (*ret_addr) &= ~BIT_31_MASK; // remove the first bit

//    printf("AddrFromPrel31: ret_addr = [%x]\n", *ret_addr);
exit_point:
    return ret;
}

//------------------------------------------------------------------------------

TBGError InitUnwindSession(elf_unwind_session *pelf) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);
    
    Elf32_Phdr phdr = {};
    CHECK_CALL(ElfFindSegment(PT_ARM_EXIDX, ELF_HEADER_BEGIN, &phdr),
               exit_point, ret);
    
    pelf->exidx = phdr.p_vaddr; 
    pelf->exidx_len = phdr.p_memsz;
    
exit_point:
    return ret;
}

//------------------------------------------------------------------------------

int FindExidxCallback(struct dl_phdr_info *info, size_t size, void *data) {
    // required address should be stored in pelf->exidx
    elf_unwind_session *pelf = (elf_unwind_session*) data;
    bool this_segment = false;
    bool has_exidx = false;
    addr_t exidx_addr = 0;
    size_t exidx_len = 0;

    int i = 0;
    for(i = 0; i < info->dlpi_phnum; ++i) {
        if(info->dlpi_phdr[i].p_type == PT_ARM_EXIDX) {
            exidx_addr = info->dlpi_phdr[i].p_vaddr;
            exidx_len = info->dlpi_phdr[i].p_memsz;
            has_exidx = true;
        }
        if(info->dlpi_phdr[i].p_type == PT_LOAD &&
           info->dlpi_addr + info->dlpi_phdr[i].p_vaddr <= pelf->exidx &&
           pelf->exidx <= info->dlpi_addr + info->dlpi_phdr[i].p_vaddr + 
                                            info->dlpi_phdr[i].p_memsz)
            this_segment = true;

        if(has_exidx && this_segment) { // found. good
            pelf->exidx = exidx_addr;
            pelf->exidx_len = exidx_len;

            return 1; // stop so-walking
        }
    }

    return 0; // to continue so-walking
}

//------------------------------------------------------------------------------

TBGError ElfFindUnwindInfo(elf_unwind_session *pelf, addr_t addr) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);
    
    pelf->exidx = addr; 
    CHECK_RET(dl_iterate_phdr(FindExidxCallback, (void*) pelf), 
              NO_PHDR_WITH_EXIDX, exit_point, ret);
    
exit_point:
    return ret;
}


//------------------------------------------------------------------------------
    
TBGError ElfGetUnwindCmd(elf_unwind_session *pelf, addr_t addr, uint8_t *buff,
                         size_t *pbuf_len) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);
    CHECK_DBG_ERR(buff, NULL_PTR);
    CHECK_DBG_ERR(pbuf_len, ZERO_LENGTH);

    addr_t exidx_entry = 0;
    CHECK_CALL(ExidxBinarySearch(pelf, addr, &exidx_entry), exit_point, ret);
    CHECK_RET(exidx_entry, NO_EXIDX_ENTRY, exit_point, ret);
    
    addr_t word_buf = 0;
    CHECK_CALL(ReadData(exidx_entry, &word_buf, 1), exit_point, ret);

//    printf("first word of exidx_entry: [%x]\n", word_buf);
    CHECK_RET(word_buf & (~BIT_31_MASK), INVALID_EXIDX_FORMAT,
              exit_point, ret); //p.15

    addr_t curr_addr = exidx_entry + sizeof(addr_t);
    CHECK_CALL(ReadData(curr_addr, &word_buf, 1), exit_point, ret);

//    printf("second word of exidx_entry: [%x]\n", word_buf);
    CHECK_RET(word_buf != EXIDX_CANTUNWIND, CANTUNWIND, exit_point, ret);

    size_t buf_cur_idx = 0;
    char add_words_count = 0; // it's char to simplify the assignation 

    // below you'll see such structure:
    // ARM_COMPACT
    // ARM_GENERIC
    //     ARM_COMPACT
    //     ARM_GENERIC

    if(EXIDX_COMPACT(word_buf)) { // EXIDX compact model
        CHECK_RET(!((word_buf >> 28) & 0x7), UNSPECIFIED_FORMAT,
                  exit_point, ret);
        int pr_index = (word_buf >> 24) & 0x0f;
        CHECK_RET(pr_index == 0, UNSPECIFIED_FORMAT, exit_point, ret);
        // read 3 bytes of data
        buff[buf_cur_idx++] = word_buf >> 16;
        buff[buf_cur_idx++] = word_buf >> 8;
        buff[buf_cur_idx++] = word_buf >> 0;
    }
    else { // EXIDX generic model
        CHECK_CALL(AddrFromPrel31(curr_addr, &curr_addr), exit_point, ret);
        CHECK_CALL(ReadData(curr_addr, &word_buf, 1), exit_point, ret);

        if(EXIDX_COMPACT(word_buf)) { // EXIDX compact model
            CHECK_RET(!((word_buf >> 28) & 0x7), INVALID_COMPACT_MODEL,
                      exit_point, ret);
            int pr_index = (word_buf >> 24) & 0x0f;
            if(pr_index == 0) { // Short format
                // Read 3 bytes of data

                buff[buf_cur_idx++] = word_buf >> 16;
                buff[buf_cur_idx++] = word_buf >> 8;
                buff[buf_cur_idx++] = word_buf >> 0;
            }
            else if(pr_index == 1 || pr_index == 2) {
                // Read number of additional 4-byte data words and the 
                // following unwind data
                add_words_count = word_buf >> 16;

                buff[buf_cur_idx++] = word_buf >> 8;
                buff[buf_cur_idx++] = word_buf >> 0;
            }
            else 
                CHECK_RET(false, UNSPECIFIED_FORMAT, exit_point, ret);
            

        }
        else { // EXIDX generic model
            curr_addr += sizeof(addr_t); // The first word is the prel31 of
                                        // personality routine. Skip it.
            CHECK_CALL(ReadData(curr_addr, &word_buf, 1), exit_point, ret);
            add_words_count = word_buf >> 24;

            buff[buf_cur_idx++] = word_buf >> 16;
            buff[buf_cur_idx++] = word_buf >> 8;
            buff[buf_cur_idx++] = word_buf >> 0;

        }
        
        CHECK_RET(0 <= add_words_count && add_words_count < 5,
                  MAX_CMD_ERROR, exit_point, ret); //5 is defined by standard
        // And now read the additional words                    
        int i = 0;
        for(i = 0; i < add_words_count; ++i) {
            CHECK_CALL(ReadData(curr_addr += sizeof(addr_t), &word_buf, 1),
                       exit_point, ret);

            buff[buf_cur_idx++] = word_buf >> 24;
            buff[buf_cur_idx++] = word_buf >> 16;
            buff[buf_cur_idx++] = word_buf >> 8;
            buff[buf_cur_idx++] = word_buf >> 0;

        }
    }

    CHECK_DBG_ERR(buf_cur_idx < MAX_CMD_BUF_LENGTH, MAX_CMD_ERROR);
    
    // Finish code is optional but in standard but needed in my algorithms
    CHECK_RET(buf_cur_idx - 1 >= 0, LOW_INDEX, exit_point, ret);
    if(buff[buf_cur_idx - 1] != ARM_FINISH_CMD_CODE)
        buff[buf_cur_idx++] = ARM_FINISH_CMD_CODE;

    (*pbuf_len) = buf_cur_idx;
exit_point:
    return ret;
}

//------------------------------------------------------------------------------

TBGError ExidxBinarySearch(elf_unwind_session *pelf, addr_t search, 
                           addr_t *pret) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pret, NULL_PTR);
    CHECK_DBG_ERR(pelf, NULL_PTR);
    CHECK_DBG_ERR(pelf->exidx, NULL_PTR);
    CHECK_DBG_ERR(pelf->exidx_len, NULL_PTR);

    addr_t cur_addr = 0;

    const addr_t sz = 2 * sizeof(addr_t); // one element size
    addr_t beg = pelf->exidx;

//    printf("pelf->exidx = [%x]\n", pelf->exidx);

    size_t min = 0;
    size_t max = pelf->exidx_len / sz - 1;
    size_t mid = 0;

//    printf("exidx table size is [%d]\n", max);

    CHECK_CALL(AddrFromPrel31(beg, &cur_addr), exit_point, ret);

    CHECK_RET(search >= cur_addr, EXIDX_NOT_CONTAIN, exit_point, ret);

//    printf("curr_addr: %[%x], search = [%x]\n", cur_addr, search);
    size_t i = 0;
    for(i = 0; i <= max; ++i) {
        CHECK_CALL(AddrFromPrel31(beg + i * sz, &cur_addr), exit_point, ret);
//        printf("exidx[%d] = [%x]\n", i, cur_addr);
    }


    CHECK_CALL(AddrFromPrel31(beg + max * sz, &cur_addr), exit_point, ret);
    if(search >= cur_addr) {
        (*pret) = beg + max * sz;
        goto exit_point;
    }

    while(max - min > 1) { // addr(max) is always bigger than search 
        mid = (min + max) >> 1;
        
        CHECK_CALL(AddrFromPrel31(mid * sz + beg, &cur_addr), exit_point, ret);
        if(search < cur_addr)
            max = mid;
        else
            min = mid;
    }

//    printf("Found at [%d] position\n", min);

    (*pret) = beg + min * sz;

    // return; // uncomment after testing
    // test the search result
    CHECK_CALL(AddrFromPrel31(beg + min * sz, &cur_addr), exit_point, ret);
    CHECK_DBG_ERR(search >= cur_addr, INVALID_BINARY_SEARCH);
    CHECK_CALL(AddrFromPrel31(beg + (min + 1) * sz, &cur_addr),
               exit_point, ret);
    CHECK_DBG_ERR(search < cur_addr, INVALID_BINARY_SEARCH);

exit_point:
    return ret;
}

